package com.digua.frame.plugin.service;

import com.lionbridge.skywalking.extension.config.Constants;
import com.lionbridge.skywalking.extension.dto.*;
import com.lionbridge.skywalking.extension.util.ComponentUtils;
import com.lionbridge.skywalking.extension.util.TraceUtils;
import org.apache.skywalking.apm.agent.core.boot.ServiceManager;
import org.apache.skywalking.apm.agent.core.commands.CommandService;
import org.apache.skywalking.apm.agent.core.conf.Config;
import org.apache.skywalking.apm.agent.core.context.IgnoredTracerContext;
import org.apache.skywalking.apm.agent.core.context.TracingContext;
import org.apache.skywalking.apm.agent.core.context.TracingContext.ListenerManager;
import org.apache.skywalking.apm.agent.core.context.trace.TraceSegment;
import org.apache.skywalking.apm.agent.core.logging.api.ILog;
import org.apache.skywalking.apm.agent.core.logging.api.LogManager;
import org.apache.skywalking.apm.agent.core.remote.GRPCChannelManager;
import org.apache.skywalking.apm.agent.core.remote.GRPCChannelStatus;
import org.apache.skywalking.apm.agent.core.remote.GRPCStreamServiceStatus;
import org.apache.skywalking.apm.agent.core.remote.TraceSegmentServiceClient;
import org.apache.skywalking.apm.commons.datacarrier.DataCarrier;
import org.apache.skywalking.apm.commons.datacarrier.buffer.BufferStrategy;
import org.apache.skywalking.apm.dependencies.com.google.gson.Gson;
import org.apache.skywalking.apm.dependencies.io.grpc.Channel;
import org.apache.skywalking.apm.dependencies.io.grpc.stub.StreamObserver;
import org.apache.skywalking.apm.network.common.v3.Commands;
import org.apache.skywalking.apm.network.common.v3.KeyStringValuePair;
import org.apache.skywalking.apm.network.language.agent.v3.*;
import org.apache.skywalking.apm.network.language.agent.v3.TraceSegmentReportServiceGrpc.TraceSegmentReportServiceStub;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.concurrent.TimeUnit;

import static org.apache.skywalking.apm.agent.core.conf.Config.Buffer.*;

/**
 * @description:
 * @author: wangzhilong
 * @date: 2021/10/20
 **/
public class TraceSegmentServiceClientExtension extends TraceSegmentServiceClient {
    private static final ILog logger = LogManager.getLogger(TraceSegmentServiceClientExtension.class);
    private static final Gson gson = new Gson();
    private long lastLogTime;
    private long segmentUplinkedCounter;
    private long segmentAbandonedCounter;
    private volatile DataCarrier<TraceSegment> carrier;
    private volatile TraceSegmentReportServiceStub serviceStub;
    private volatile GRPCChannelStatus status;

    public TraceSegmentServiceClientExtension() {
        this.status = GRPCChannelStatus.DISCONNECT;
    }

    @Override
    public void prepare() {
        if (Constants.KEEP_REPORT) {
            ((GRPCChannelManager) ServiceManager.INSTANCE.findService(GRPCChannelManager.class)).addChannelListener(this);
        }

    }

    @Override
    public void boot() {
        if (Constants.KEEP_REPORT) {
            this.lastLogTime = System.currentTimeMillis();
            this.segmentUplinkedCounter = 0L;
            this.segmentAbandonedCounter = 0L;
            this.carrier = new DataCarrier(CHANNEL_SIZE, BUFFER_SIZE);
            this.carrier.setBufferStrategy(BufferStrategy.IF_POSSIBLE);
            this.carrier.consume(this, 1);
        }

    }

    @Override
    public void onComplete() {
        ListenerManager.add(this);
    }

    @Override
    public void shutdown() {
        ListenerManager.remove(this);
        if (Constants.KEEP_REPORT) {
            this.carrier.shutdownConsumers();
        }

    }

    @Override
    public void init() {
    }

    @Override
    public void consume(List<TraceSegment> data) {
        if (GRPCChannelStatus.CONNECTED.equals(this.status)) {
            final GRPCStreamServiceStatus status = new GRPCStreamServiceStatus(false);
            StreamObserver upstreamSegmentStreamObserver = ((TraceSegmentReportServiceStub)this.serviceStub.withDeadlineAfter((long) Config.Collector.GRPC_UPSTREAM_TIMEOUT, TimeUnit.SECONDS)).collect(new StreamObserver<Commands>() {
                @Override
                public void onNext(Commands commands) {
                    ((CommandService)ServiceManager.INSTANCE.findService(CommandService.class)).receiveCommand(commands);
                }

                @Override
                public void onError(Throwable throwable) {
                    status.finished();
                    if (TraceSegmentServiceClientExtension.logger.isErrorEnable()) {
                        TraceSegmentServiceClientExtension.logger.error(throwable, "Send UpstreamSegment to collector fail with a grpc internal exception.", new Object[0]);
                    }

                    ((GRPCChannelManager)ServiceManager.INSTANCE.findService(GRPCChannelManager.class)).reportError(throwable);
                }

                @Override
                public void onCompleted() {
                    status.finished();
                }
            });

            try {
                SegmentObject upstreamSegment = null;
                Iterator var5 = data.iterator();

                while(var5.hasNext()) {
                    TraceSegment segment = (TraceSegment)var5.next();
                    upstreamSegment = segment.transform();
                    upstreamSegmentStreamObserver.onNext(upstreamSegment);
                    this.segmentTransform(upstreamSegment);
                }
            } catch (Throwable var7) {
                logger.error(var7, "Transform and send UpstreamSegment to collector fail.", new Object[0]);
            }

            upstreamSegmentStreamObserver.onCompleted();
            status.wait4Finish();
            this.segmentUplinkedCounter += (long)data.size();
        } else {
            this.segmentAbandonedCounter += (long)data.size();
        }

        this.printUplinkStatus();
    }

    private void printUplinkStatus() {
        long currentTimeMillis = System.currentTimeMillis();
        if (currentTimeMillis - this.lastLogTime > 30000L) {
            this.lastLogTime = currentTimeMillis;
            if (this.segmentUplinkedCounter > 0L) {
                logger.debug("{} trace segments have been sent to collector.", new Object[]{this.segmentUplinkedCounter});
                this.segmentUplinkedCounter = 0L;
            }

            if (this.segmentAbandonedCounter > 0L) {
                logger.debug("{} trace segments have been abandoned, cause by no available channel.", new Object[]{this.segmentAbandonedCounter});
                this.segmentAbandonedCounter = 0L;
            }
        }

    }

    @Override
    public void onError(List<TraceSegment> data, Throwable t) {
        logger.error(t, "Try to send {} trace segments to collector, with unexpected exception.", new Object[]{data.size()});
    }

    @Override
    public void onExit() {
    }

    @Override
    public void afterFinished(TraceSegment traceSegment) {
        if (Constants.KEEP_REPORT) {
            if (traceSegment.isIgnore()) {
                return;
            }

            if (!this.carrier.produce(traceSegment) && logger.isDebugEnable()) {
                logger.debug("One trace segment has been abandoned, cause by buffer is full.");
            }
        } else {
            try {
                this.segmentTransform(traceSegment.transform());
            } catch (Throwable var3) {
                logger.error(var3, "Transform and send UpstreamSegment to collector fail.", new Object[0]);
            }
        }

    }

    @Override
    public void statusChanged(GRPCChannelStatus status) {
        if (Constants.KEEP_REPORT) {
            if (GRPCChannelStatus.CONNECTED.equals(status)) {
                Channel channel = ((GRPCChannelManager)ServiceManager.INSTANCE.findService(GRPCChannelManager.class)).getChannel();
                this.serviceStub = TraceSegmentReportServiceGrpc.newStub(channel);
            }

            this.status = status;
        }

    }

    private void segmentTransform(SegmentObject segmentObject) {
        SegmentDTO segment = new SegmentDTO();
        segment.setHost(TraceUtils.getHost());
        segment.setTraceId(segmentObject.getTraceId());
        segment.setSegmentId(segmentObject.getTraceSegmentId());
        segment.setServiceName(TraceUtils.getServiceName());
        segment.setServiceInstanceName(TraceUtils.getServiceInstanceName());
        segment.setSizeLimited(segmentObject.getIsSizeLimited());
        List<SpanObject> spanObjects = segmentObject.getSpansList();
        if (null != spanObjects && spanObjects.size() > 0) {
            List<SpanDTO> spans = new ArrayList();
            boolean error = false;
            SpanDTO span = null;
            Iterator var7 = spanObjects.iterator();

            while(var7.hasNext()) {
                SpanObject spanObject = (SpanObject)var7.next();
                span = this.spanTransform(spanObject, segment);
                if (null != span) {
                    if (span.isError()) {
                        error = true;
                    }

                    spans.add(span);
                }
            }

            segment.setSpans(spans);
            segment.setError(error);
            if (segment.isError() && spans.size() == 1 && span.getSpanType().equals("Entry") && (span.getComponentId() == 1 || span.getComponentId() == 19 || span.getComponentId() == 84)) {
                segment.setSkipAnalysis(true);
            }
        }

        logger.info(gson.toJson(segment));
    }

    private SpanDTO spanTransform(SpanObject spanObject, SegmentDTO segment) {
        SpanDTO span = new SpanDTO();
        if (null != spanObject) {
            span = new SpanDTO();
            span.setSpanId(spanObject.getSpanId());
            span.setParentSpanId(spanObject.getParentSpanId());
            if (spanObject.getSpanId() == 0) {
                segment.setStartTime(spanObject.getStartTime());
                segment.setEndTime(spanObject.getEndTime());
            }

            span.setStartTime(spanObject.getStartTime());
            span.setEndTime(spanObject.getEndTime());
            span.setOperationName(spanObject.getOperationName());
            span.setComponentId(spanObject.getComponentId());
            span.setComponentName(ComponentUtils.getComponentName(spanObject.getComponentId()));
            if (spanObject.getSpanTypeValue() == 0) {
                span.setSpanType("Entry");
            } else if (spanObject.getSpanTypeValue() == 1) {
                span.setSpanType("Exit");
            } else {
                span.setSpanType("Local");
            }

            if (null != spanObject.getSpanLayer()) {
                span.setSpanLayer(spanObject.getSpanLayer().toString());
            }

            span.setError(spanObject.getIsError());
            List<KeyStringValuePair> tagKvPairs = spanObject.getTagsList();
            Iterator iterator;
            KeyStringValuePair log;
            if (null != tagKvPairs && tagKvPairs.size() > 0) {
                List<KvDTO> tags = new ArrayList();
                KvDTO kv = null;
                iterator = tagKvPairs.iterator();

                while(iterator.hasNext()) {
                    log = (KeyStringValuePair)iterator.next();
                    kv = new KvDTO();
                    kv.setKey(log.getKey());
                    kv.setValue(log.getValue());
                    tags.add(kv);
                }

                span.setTags(tags);
            }

            List<SegmentReference> refsList = spanObject.getRefsList();
            SegmentReference key;
            if (null != refsList && refsList.size() > 0) {
                List<SegmentReferenceDTO> refs = new ArrayList();
                iterator = null;
                Iterator var20 = refsList.iterator();

                while(var20.hasNext()) {
                    key = (SegmentReference)var20.next();
                    SegmentReferenceDTO segmentReferenceDTO = new SegmentReferenceDTO();
                    segmentReferenceDTO.setParentTraceSegmentId(key.getParentTraceSegmentId());
                    segmentReferenceDTO.setRefType(key.getRefType().toString());
                    segmentReferenceDTO.setParentSpanId(key.getParentSpanId());
                    segmentReferenceDTO.setParentServiceName(key.getParentService());
                    segmentReferenceDTO.setParentServiceInstanceName(key.getParentServiceInstance());
                    segmentReferenceDTO.setParentEndpoint(key.getParentEndpoint());
                    segmentReferenceDTO.setNetworkAddressUsedAtPeer(key.getNetworkAddressUsedAtPeer());
                    refs.add(segmentReferenceDTO);
                }

                span.setRefs(refs);
            }

            List<Log> logsList = spanObject.getLogsList();
            if (null != logsList && logsList.size() > 0) {
                List<LogDTO> logs = new ArrayList();
                log = null;
                key = null;
                List<KeyStringValuePair> kvPairs = null;
                Iterator var11 = logsList.iterator();

                while(true) {
                    do {
                        do {
                            if (!var11.hasNext()) {
                                span.setLogs(logs);
                                return span;
                            }

                            Log logObject = (Log)var11.next();
                            kvPairs = logObject.getDataList();
                        } while(null == kvPairs);
                    } while(kvPairs.size() <= 0);

                    LogDTO logDTO = new LogDTO();
                    Iterator var13 = kvPairs.iterator();

                    while(var13.hasNext()) {
                        KeyStringValuePair kvPair = (KeyStringValuePair)var13.next();
                        String pairKey = kvPair.getKey();
                        if ("event".equals(pairKey)) {
                            logDTO.setEvent(kvPair.getValue());
                        } else if ("error.kind".equals(pairKey)) {
                            logDTO.setErrorKind(kvPair.getValue());
                        } else if ("message".equals(pairKey)) {
                            logDTO.setMessage(kvPair.getValue());
                        } else if ("stack".equals(pairKey)) {
                            logDTO.setStack(kvPair.getValue());
                        }
                    }

                    logs.add(logDTO);
                }
            }
        }

        return span;
    }

}
